package LibDL.eval.regression;

import LibDL.core.Scalar;
import LibDL.core.Tensor;
import LibDL.core.functional;

import java.util.Arrays;

public class RMSEEvaluator extends AbstractRegressionEvaluator{
    @Override
    protected double core(Tensor truth, Tensor pred) {
        TargetType type = checkTarget(truth,pred);
        if (!Arrays.asList(TargetType.SINGLE_OUTPUT, TargetType.MULTI_OUTPUT).contains(type)){
            throw new IllegalArgumentException("The type of input is not supported." +
                    " Found: "+type);
        }
        if (getSampleWeight() == null) {
            setSampleWeight(functional.ones(truth.size(0)));
        }
        if (type.equals(TargetType.MULTI_OUTPUT) && getOutputWeight() == null) {
            setOutputWeight(functional.ones(truth.size(1)));
        }
        checkConsistentShape(truth,getSampleWeight(),getOutputWeight(),type);

        Tensor diff = truth.sub(pred);
        Tensor evalVal;
        evalVal = weightedSum(diff.mul(diff),getSampleWeight(),true,0);
        if (type.equals(TargetType.MULTI_OUTPUT)){
            evalVal = weightedSum(evalVal,getOutputWeight(),true,0);
        }
        evalVal = evalVal.pow(new Scalar(0.5));
        return evalVal.item().to_double();
    }
}
